package com.katze.ssh.cmd;

import com.katze.ssh.IChannel;
import com.katze.ssh.SSHException;
import com.katze.ssh.SshSessionEntry;
import org.apache.sshd.client.channel.ChannelExec;
import org.apache.sshd.client.channel.ClientChannelEvent;
import org.apache.sshd.client.session.ClientSession;
import org.apache.sshd.common.SshException;
import org.apache.sshd.common.channel.exception.SshChannelOpenException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;
import java.util.EnumSet;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.function.Consumer;

public class ExecChannel extends ExecCommand implements IChannel {
    private static final Logger log = LoggerFactory.getLogger(ExecChannel.class);

    private final SshSessionEntry session;
    private ClientSession clientSession;

    public ExecChannel(SshSessionEntry session) {
        this.session = session;
    }

    @Override
    public void open() throws Exception{
        this.clientSession = session.getClientSession();
    }

    public void script(String position, String script, Consumer<String> action) throws Exception {
//        source /etc/profile;source ~/.bash_profile;source ~/.bashrc;
        String base = "source ~/.bash_profile; cd ${position} && ${script}";
        Map<String, String> variables = new HashMap<>();
        variables.put("position", position);
        variables.put("script", script);
        execute(toCommand(base, variables), action);
    }

    public void kill(String pid, Consumer<String> action) throws Exception {
        String base = "source ~/.bash_profile; kill ${pid}";
        Map<String, String> variables = new HashMap<>();
        variables.put("pid", pid);
        execute(toCommand(base, variables), action);
    }

    /**
     * 执行ps -ef命令
     * @UID (User ID):
     * 运行该进程的用户的标识符。可以是实际的用户 ID 或者用户名。
     * @PID (Process ID):
     * 进程的唯一标识符。每个进程都有一个唯一的 PID。
     * @PPID (Parent Process ID):
     * 创建此进程的父进程的 PID。如果父进程已经结束，那么该值通常是 1（init 进程或 systemd）。
     * @CPU (CPU Usage):
     * 进程使用的 CPU 时间比例，自进程启动以来累计计算。它表示的是百分比，但不是实时的 CPU 使用率。
     * @STIME (Start Time):
     * 进程开始的时间。对于较早启动的进程，这可能是日期；对于最近启动的进程，则是当天的时间。
     * @TTY (Controlling Terminal):
     * 分配给该进程的终端（tty）。如果没有关联的终端，通常显示为 ?。
     * @TIME (CPU Time):
     * 进程累积使用 CPU 的总时间，格式为 [dd-]hh:mm:ss。这里的时间仅指 CPU 执行时间，而非进程运行的全部时间。
     * @CMD (Command with Arguments):
     * 启动该进程时所用的命令及其参数。对于某些服务或守护程序，可能会显示可执行文件路径或者服务名称。
     * @param conditions grep 条件
     */
    public void  psEf(String condition, Consumer<Properties> action) throws Exception {
       psEf(new String[]{condition}, action);
    }

    public void  psEf(String[] conditions, Consumer<Properties> action) throws Exception {
        String base = "ps -ef ${conditions}|grep " + session.getUsername() + " |grep -v 'grep'";
        Map<String, String> variables = new HashMap<>();
        for (String o : conditions) {
            // 33是'!'的ASCII码
            if (o.charAt(0) == 33) {
                o = "-v " + o.substring(1); // 将"!xxx"转换为"-v xxx"（反向过滤）
            }
            variables.merge("conditions", "|grep " + o + " ", (v1, v2) -> v1 + v2);
        }
        execute(toCommand(base, variables), 8, action);
    }

    /**
     * 执行ps -u xxx -o pid,etimes=times,%cpu,%mem,cmd命令
     * @param columns 指定需要显示的列名
     */
    public void psUO(Consumer<Properties> action) throws Exception {
        String base = "ps -u ${username} -o ${content} --sort=etimes |grep -Ev 'ps -u ${username}|sshd:'";
        Map<String, String> variables = new HashMap<>();
        String content = "pid=pid,etimes=times,%cpu=cpuUsage,%mem=memUsage,cmd=cmd";
        boolean isExist = false;
//        if (columns.length > 0) {
//            StringBuilder temp = new StringBuilder();
//            for (String column : columns) {
//                if (!isExist) {
//                    isExist = content.contains(column);
//                }
//                temp.append(",").append(column);
//            }
//            if (isExist) {
//                content = temp.deleteCharAt(0).toString();
//            } else {
//                content += temp.toString();
//            }
//        }
        variables.put("content", content);
        variables.put("username", session.getUsername());
        execute(toCommand(base, variables), 0, action);
    }

    public void  execute(String cmd, int columns, Consumer<Properties> action) throws Exception {
        try (ChannelExec channel = clientSession.createExecChannel(cmd); ByteArrayOutputStream err = new ByteArrayOutputStream()) {
            channel.setErr(err);
            channel.open().verify(5, TimeUnit.SECONDS);
            channel.waitFor(EnumSet.of(ClientChannelEvent.CLOSED), 0);
            String msg = err.toString();
            if (msg.isBlank()) {
                 resolve(channel.getInvertedOut(), columns, action);
            } else {
                throw new SSHException(msg);
            }
        } catch (SshException e) {
            if (e.getCause() instanceof TimeoutException) {
                throw new SSHException("会话连接超时");
            } if (e.getCause() instanceof SshChannelOpenException) {
                session.getClientSession().close();
                throw new SSHException("会话连接异常，请重试");
            } else {
                throw e;
            }
        }
    }

    public void execute(String cmd, Consumer<String> action) throws Exception{
        log.info("session: {}@{}:{}, execute cmd: {}", session.getUsername(), session.getHost(), session.getPort(), cmd);
        try (ChannelExec channel = clientSession.createExecChannel(cmd); ByteArrayOutputStream out = new ByteArrayOutputStream(); ByteArrayOutputStream err = new ByteArrayOutputStream()) {
            channel.setOut(out);
            channel.setErr(err);
            channel.open().verify(5, TimeUnit.SECONDS);
            channel.waitFor(EnumSet.of(ClientChannelEvent.CLOSED), 0);
            channel.close(false); //true:立马关闭；false:优雅关闭
            String msg = err.toString();
            if (msg.isBlank()) {
                action.accept(out.toString(StandardCharsets.UTF_8));
            } else {
                throw new IllegalStateException(msg);
            }
        } catch (SshException e) {
            if (e.getCause() instanceof TimeoutException) {
                throw new SSHException("会话连接超时");
            } if (e.getCause() instanceof SshChannelOpenException) {
                session.getClientSession().close();
                throw new SSHException("会话连接异常，请重试");
            } else {
                throw e;
            }
        }
    }
}
